#include "flash_tool.h"
#include "md5.h"
#include <string.h> //strlen
#include <stdio.h> //printf
#include <stdlib.h> //system

static int calc_file_md5sum(int fd, off_t size, uint8_t md5sum[16])
{
    int ret = 0;
    MD5_CTX md5_ctx;
    off_t left = size;
    const uint32_t step = 1024 * 1024;
    uint8_t buffer[step];
    ssize_t block_size;
    MD5Init(&md5_ctx);

    while(left > 0)
    {
        if(left > step)
        {
            block_size = step;
        }
        else
        {
            block_size = left;
        }
        ssize_t got = read(fd, buffer, block_size);
        if(got != block_size)
        {
            ret = -EIO;
            break;
        }
        MD5Update(&md5_ctx, buffer, block_size);
        left -= block_size;
    }

    MD5Final(&md5_ctx, md5sum);
    return ret;
}

static int copy_file(int dst_fd, int src_fd, off_t size)
{
    off_t left = size;
    const ssize_t step = 1024 * 1024;
    char buffer[step];
    ssize_t cur_step;
    ssize_t got;
    int persent = 0;
    int err = 0;

    while(left > 0 && err == 0)
    {
        if(left > step)
        {
            cur_step = step;
        }
        else
        {
            cur_step = left;
        }
        got = read(src_fd, buffer, cur_step);
        if(got == cur_step)
        {
            got = write(dst_fd, buffer, cur_step);
            if(got != cur_step)
            {
                err = -EIO;
                break;
            }
        }
        else
        {
            err = -EIO;
            break;
        }
        left -= cur_step;
        int p = 100 - ((left * 100) / size);
        if(p != persent)
        {
            persent = p;
            printf("\r %10lld %10lld (%2d%% done)", (long long int)(size - left), (long long int)size, persent);
            fflush(stdout);
            syncfs(dst_fd);//每完成1％数据复制即下盘
        }
    }
    return err;
}

static int check_section(int fd, off_t& data_left)
{
    if((size_t)data_left > sizeof(SectionHeader))
    {
        int ret = 0;
        SectionHeader header;
        if(read(fd, (void*)&header, sizeof(header)) != sizeof(header))
        {
            return -EIO;
        }
        if(data_left < (off_t)(header.data_size + sizeof(header)))
        {
            return -1;
        }
        unsigned char md5_bak[sizeof(header.md5)];
        memcpy(md5_bak, header.md5, sizeof(md5_bak));//备份原值
        memset(&header.md5, 0, sizeof(header.md5));//清空

        MD5_CTX md5_ctx;
        unsigned char md5_sum[16];
        off_t left = header.data_size;
        const uint32_t step = 1024 * 1024;
        uint8_t buffer[step];
        ssize_t block_size;
        MD5Init(&md5_ctx);
        MD5Update(&md5_ctx, (unsigned char*)&header, sizeof(header));

        while(left > 0)
        {
            if(left > step)
            {
                block_size = step;
            }
            else
            {
                block_size = left;
            }
            ssize_t got = read(fd, buffer, block_size);
            if(got != block_size)
            {
                ret = -EIO;
                break;
            }
            MD5Update(&md5_ctx, buffer, block_size);
            left -= block_size;
        }

        MD5Final(&md5_ctx, md5_sum);

        //比较md5值
        for(size_t idx = 0; idx < sizeof(md5_sum) && ret == 0; ++idx)
        {
            if(md5_sum[idx] != md5_bak[idx])
            {
                ret = -1;
                break;
            }
        }
        if(ret == 0)
        {
            data_left -= (sizeof(header) + header.data_size);
        }
        return ret;
    }
    return -1;
}

int check_binary_file(const char* path)
{
    int ret = 0;
    if(path != NULL)
    {
        int fd = open(path, O_RDONLY);
        if(fd < 0)
        {
            return -errno;
        }
        off_t file_size = lseek(fd, 0, SEEK_END);
        if(file_size > (off_t)sizeof(SectionHeader))
        {
            lseek(fd, 0, SEEK_SET);
            off_t data_left = file_size;
            ret = 0;

            while(ret == 0 && data_left > 0)
            {
                ret = check_section(fd, data_left);
            }
            if(data_left > 0)
            {
                ret = -1;
            }
        }
        else
        {
            ret = -1;
        }

        close(fd);

        return ret;
    }
    return -1;
}

static int execute_command(int src_fd, const SectionHeader& header, off_t& data_left, const off_t total, int report_fd)
{
    //两个字节长度，后跟命令字符串
    if(header.data_size <= 2 || header.data_size > 4095)
    {
        return -1;
    }
    uint16_t cmd_flag = 0;
    uint16_t cmd_size;
    char cmd[4096];
    if(read(src_fd, &cmd_flag, 2) != 2)
    {
        return -EIO;
    }
    if(read(src_fd, &cmd_size, 2) != 2)
    {
        return -EIO;
    }
    if(header.data_size != (uint64_t)(cmd_size + 4))
    {
        return -1;
    }
    if(read(src_fd, cmd, cmd_size) != cmd_size)
    {
        return -EIO;
    }
    if(cmd_flag & 1)
    {
        printf("will execute command: %s\n", cmd);
    }
    system(cmd);//不关心这个返回值
    return 0;

}

static int flash_fs_file(int src_fd, const SectionHeader& header, off_t& data_left, const off_t total, int report_fd)
{
    //2字节分区号，2字节路径长度，路径文件数据
    if(header.data_size < 4)
    {
        return -1;
    }
    char info[260];
    if(read(src_fd, info, 4) != 4)
    {
        return -EIO;
    }
    uint16_t path_len = *(uint16_t*)(info + 2);
    if(path_len > 255)
    {
        return -1;
    }
    if(read(src_fd, info + 4, path_len) != path_len)
    {
        return -1;
    }
    info[path_len + 4] = 0;
    data_left -= (path_len + 4);

    int dst_file = open(info + 4, O_WRONLY|O_CREAT, S_IRWXU|S_IXUSR|S_IRWXG);
    if(dst_file <= 0)
    {
        dst_file = open(info + 4, O_WRONLY|O_CREAT);
    }
    if(dst_file > 0)
    {
        data_left -= (header.data_size + sizeof(header));
        int err = copy_file(dst_file, src_fd, header.data_size - (path_len + 4));
        if(err == 0)
        {
            char cmd[512];
            snprintf(cmd, 512, "chmod +xrw %s", info + 4);
            system(cmd);
            return err;
        }
        return 0;
    }
    else
    {
        return -errno;
    }
}

static int flash_section(int src_fd, int dst_fd, off_t& data_left, const off_t total, int report_fd)
{
    SectionHeader header;
    ssize_t got = read(src_fd, &header, sizeof(header));

    if(got == sizeof(header))
    {
        if(header.flag == SECTION_RAW_DATA)
        {
            int ret = 0;
            uint64_t left = header.data_size;
            const ssize_t step = 1024 * 1024;
            ssize_t cur_step;
            char buffer[step];
            int persent = 0;
            char msg[128];

            off_t pos = lseek(dst_fd, header.blk_offset, SEEK_SET);
            printf("start flashing at %lld (%llu), size:%llu...\n", (long long int)pos, (long long unsigned int)header.blk_offset, (long long unsigned)header.data_size);


            while(left > 0 && ret == 0)
            {
                if(left > (off_t)step)
                {
                    cur_step = step;
                }
                else
                {
                    cur_step = left;
                }
                got = read(src_fd, buffer, cur_step);
                if(got == cur_step)
                {
                    got = write(dst_fd, buffer, cur_step);
                    if(got != cur_step)
                    {
                        ret = -EIO;
                        break;
                    }
                }
                else
                {
                    ret = -EIO;
                    break;
                }
                left -= cur_step;
                float p = 100 - ((data_left - (header.data_size - left)) * 100.0) / total;

                printf("\r %0.2f persent done", p);
                if((int)p % 5 == 0)
                {
                    syncfs(dst_fd);
                    //if(report_fd > 0)
                    {
                        persent = p;
                        snprintf(msg, 128, "echo \"PROGRESS %d\" >/mnt/.psplash/psplash_fifo && sync", persent);
                        //printf("%s\n", msg);
                        system(msg);
                        sleep(1);
                        //printf("%s\n", msg);
                        //write(report_fd, msg, strlen(msg) + 1);
                        //syncfs(report_fd);
                    }
                }
                fflush(stdout);
                persent = p;
                //usleep(50000);
            }

            if(ret == 0)
            {
                data_left -= (header.data_size + sizeof(header));
            }
            return ret;
        }
        else if(header.flag == SECTION_FS_FILE)
        {
            //文件系统文件
            return flash_fs_file(src_fd, header, data_left, total, report_fd);
        }
        else if(header.flag == SECTION_COMMAND)
        {
            //执行本地命令
            return execute_command(src_fd, header, data_left, total, report_fd);
        }
    }
    return -1;
}

int flash_binary(const char* file, const char* device, const char* report_file)
{
    int ret = 0;
    int fi, fo;

    if(file == NULL || device == NULL)
    {
        return -1;
    }

    if(access(device, F_OK) == 0)
    {
        fo = open(device, O_WRONLY|O_TRUNC);
    }
    else
    {
        fo = open(device, O_WRONLY|O_CREAT|O_TRUNC,  S_IRWXU|S_IXUSR|S_IRWXG);
    }
    
    fi = open(file, O_RDONLY);

    if(fi > 0 && fo > 0)
    {
        //文件已经检验过，所以不做错误检查
        off_t file_size = lseek(fi, 0, SEEK_END);
        lseek(fi, 0, SEEK_SET);

        int report_fd = 0;
        if(report_file != NULL)
        {
            report_fd = open(report_file, O_WRONLY|O_NONBLOCK);
        }

        off_t file_left = file_size;
        while(file_left > 0 && ret == 0)
        {
            ret = flash_section(fi, fo, file_left, file_size, report_fd);
        }

        if(ret == 0)
        {
            off_t boot_flag_pos = 2621400;
            lseek(fo, boot_flag_pos, SEEK_SET);
            uint32_t normal_boot_flag = 0;
            write(fo, &normal_boot_flag, sizeof(normal_boot_flag));
        }
    }
    else
    {
        ret = -errno;
    }
    if(fi > 0)
    {
        close(fi);
    }
    if(fo > 0)
    {
        close(fo);
    }
    return ret;
}

int build_binary(int argc, char* argv[])
{
    return 0;
}

int build_section(const char* binary_file, const char* desc, const char* dst_file, uint64_t flash_offt, const char* dst_fs_path, const char* cmd, uint64_t source_offset, uint64_t max_source_size)
{
    if((binary_file != NULL || cmd != NULL)&& dst_file != NULL)
    {
        int fi = -1, fo;
        int err = 0;

        if(cmd == NULL)
        {
        fi = open(binary_file, O_RDONLY);
        if(fi <= 0)
        {
            return -EIO;
        }
        }

        fo = open(dst_file, O_RDWR|O_CREAT, S_IRUSR|S_IWUSR|S_IRGRP);
        //if(fo < 0 && )

        if( fo > 0)
        {
            ssize_t file_info_len = 0;
            SectionHeader header;
            memset(&header, 0, sizeof(header));
            header.blk_offset = flash_offt;
            off_t src_file_size =  lseek(fi, 0, SEEK_END);
            if((src_file_size - source_offset - max_source_size) >= 0)
            {
                if(max_source_size > 0)
                    header.data_size = max_source_size;
                else
                    header.data_size = src_file_size - source_offset;
            }
            else
            {
                //source data file is too small to fill destination as parameter descripted
                close(fi);
                close(fo);
                return false;
            }
            lseek(fi, source_offset, SEEK_SET);

            if(dst_fs_path != NULL)
            {
                header.flag = SECTION_FS_FILE;
                file_info_len = 4 + strlen(dst_fs_path) + 1;
                header.data_size += file_info_len;
            }
            else if(cmd != NULL)
            {
                header.flag = SECTION_COMMAND;
                header.blk_offset = 0;
                header.data_size = (strlen(cmd) + 1 + 2 + 2);
            }

            //定位到输出文件末尾写
            off_t fo_pos = lseek(fo, 0, SEEK_END);
            ssize_t got = write(fo, &header, sizeof(header));
            if(got == sizeof(header))
            {

                if(dst_fs_path != NULL)
                {
                    //是一个文件类型数据
                    char info[260];
                    *(uint16_t*)info = 0;
                    *(uint16_t*)(info + 2) = file_info_len - 4;
                    strncpy((char*)info + 4, dst_fs_path, 255);
                    if(write(fo, info, file_info_len) != file_info_len)
                    {
                        return -EIO;
                    }
                }
                if(header.flag == SECTION_COMMAND)
                {
                    uint16_t cmd_len = strlen(cmd) + 1;
                    uint16_t flag = 1;
                    if(write(fo, &flag, 2) != 2)
                    {
                        return -EIO;
                    }
                    if(write(fo, &cmd_len, 2) != 2)
                    {
                        return -EIO;
                    }
                    if(write(fo, cmd, cmd_len) != cmd_len)
                    {
                        return -1;
                    }
                }
                else if(copy_file(fo, fi, header.data_size - file_info_len) != 0)
                {
                    return -EIO;
                }
                //写数据完成
                lseek(fo, fo_pos, SEEK_SET);
                //计算md5
                uint8_t md5sum[16];
                calc_file_md5sum(fo, header.data_size + sizeof(header), md5sum);
                lseek(fo, fo_pos + ((char*)&header.md5 - (char*)&header), SEEK_SET);

                got = write(fo, md5sum, sizeof(md5sum));
                if(got != sizeof(md5sum))
                {
                    err = -EIO;
                }

            }
        }
        else
        {
            err = -EIO;
        }
        if(fi > 0)
        {
            close(fi);
        }
        if(fo > 0)
        {
            close(fo);
        }
        return err;
    }
    return -1;
}
